import sys
import akshare as ak
import pyqtgraph as pg
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, 
                            QWidget, QLineEdit, QPushButton, QHBoxLayout)

class StockKLineChart(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('股票K线图')
        self.setGeometry(100, 100, 800, 600)
        
        # 创建主部件和布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)
        
        # 创建输入框和按钮
        input_layout = QHBoxLayout()
        self.stock_code_input = QLineEdit()
        self.stock_code_input.setPlaceholderText('输入股票代码(如: sh600000)')
        self.plot_button = QPushButton('确定')
        self.plot_button.clicked.connect(self.plot_kline)
        
        input_layout.addWidget(self.stock_code_input)
        input_layout.addWidget(self.plot_button)
        layout.addLayout(input_layout)
        
        # 创建绘图区域
        self.plot_widget = pg.PlotWidget()
        layout.addWidget(self.plot_widget)
        
    def plot_kline(self):
        stock_code = self.stock_code_input.text().strip()
        if not stock_code:
            return
            
        try:
            # 获取股票数据
            df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", adjust="hfq")
            print(f"获取到数据条数: {len(df)}")  # 调试信息
            if df.empty:
                print("获取的数据为空，请检查股票代码是否正确")
                return
                
            df['date'] = df['日期'].apply(pg.toDate)
            
            # 清除之前的绘图
            self.plot_widget.clear()
            
            # 创建K线图
            kline = pg.CandlestickItem(
                x=df['date'],
                open=df['开盘'],
                high=df['最高'],
                low=df['最低'],
                close=df['收盘']
            )
            
            self.plot_widget.addItem(kline)
            self.plot_widget.setLabel('left', '价格')
            self.plot_widget.setLabel('bottom', '日期')
            self.plot_widget.showGrid(x=True, y=True)
            
            # 设置合适的显示范围
            self.plot_widget.setXRange(min(df['date']), max(df['date']))
            y_min = min(df['最低']) * 0.98
            y_max = max(df['最高']) * 1.02
            self.plot_widget.setYRange(y_min, y_max)
            
            # 启用缩放功能
            self.plot_widget.setMouseEnabled(x=True, y=True)
            
        except Exception as e:
            print(f"获取数据失败: {e}")
            import traceback
            traceback.print_exc()  # 打印完整的错误堆栈

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = StockKLineChart()
    window.show()
    sys.exit(app.exec_())
